package middleware

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"github.com/gin-gonic/gin"
	"net/http"
	"strconv"
	"time"
)

// Config AK/SK 中间件配置
type Config struct {
	// 获取 SecretKey 的函数类型
	GetSecretKey func(accessKey string) (string, error)
	// 签名有效期（秒）
	Expires int64
	// 请求头字段名称（可自定义）
	HeaderAccessKey string
	HeaderSignature string
	HeaderTimestamp string
}

// DefaultConfig 默认配置
func DefaultConfig(getSecretKey func(string) (string, error)) *Config {
	return &Config{
		GetSecretKey:     getSecretKey,
		Expires:          300,
		HeaderAccessKey:  "X-Access-Key",
		HeaderSignature:  "X-Signature",
		HeaderTimestamp:  "X-Timestamp",
	}
}

// AK/SK 验证中间件
func AKSKMiddleware(config *Config) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 获取请求头
		accessKey := c.GetHeader(config.HeaderAccessKey)
		signature := c.GetHeader(config.HeaderSignature)
		timestampStr := c.GetHeader(config.HeaderTimestamp)

		// 检查必要头是否存在
		if accessKey == "" || signature == "" || timestampStr == "" {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"message": "Missing authentication headers",
			})
			return
		}

		// 解析时间戳
		timestamp, err := strconv.ParseInt(timestampStr, 10, 64)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"message": "Invalid timestamp format",
			})
			return
		}

		// 检查时间有效期
		now := time.Now().Unix()
		if timestamp < now-config.Expires || timestamp > now+config.Expires {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"message": "Request expired",
			})
			return
		}

		// 获取 SecretKey
		secretKey, err := config.GetSecretKey(accessKey)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"message": "Invalid access key",
			})
			return
		}

		// 生成签名
		expectedSignature := GenerateSignature(
			c.Request.Method,
			c.Request.URL.Path,
			timestampStr,
			secretKey,
		)

		// 验证签名
		if !hmac.Equal([]byte(signature), []byte(expectedSignature)) {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"message": "Invalid signature",
			})
			return
		}

		c.Next()
	}
}

// 生成签名（客户端和服务端要保持一致）
func GenerateSignature(method, path, timestamp, secretKey string) string {
	payload := fmt.Sprintf("%s\n%s\n%s", method, path, timestamp)
	h := hmac.New(sha256.New, []byte(secretKey))
	h.Write([]byte(payload))
	return hex.EncodeToString(h.Sum(nil))
}